import os
import sys
import numpy as np
import matplotlib.pyplot as plt

# cmd_shell = "grep 'loss:' worker_31/work.log | grep '\[' | awk -F 'loss:' '{print $2}'| " \
#             "awk -F '\[' '{print $2}' | awk -F '/' '{print $1}' > loss.txt"
# os.system(cmd_shell)
if len(sys.argv) > 1:
    loss_files = []
    for i in range(1, 10):
        try:
            loss_files.append(sys.argv[i])
        except:
            pass
else:
    loss_path = "./"
    if os.path.isdir(loss_path):
        loss_files = sorted([f for f in os.listdir(loss_path) if f.startswith("loss") and f.endswith(".txt")])

loss_data = {}
for i, file in enumerate(loss_files):
    filename = os.path.splitext(file)[0]
    loss = np.loadtxt(file)
    loss_data[filename] = loss

labels = list(loss_data.keys())
min_data_len = min([len(loss) for loss in list(loss_data.values())])
data_plot = [loss[:min_data_len] for loss in list(loss_data.values())]
axis = np.arange(0, min_data_len)
for d in data_plot:
    plt.plot(axis, d)
plt.legend(labels=labels)
plt.xlabel("step")
plt.ylabel("loss")
plt.savefig(f'./loss.png', dpi=1000, bbox_inches='tight')
